import torch
import torch.nn.functional as F
import numpy as np
from PIL import Image
from tqdm import tqdm
import os
# 设置非交互式matplotlib后端以避免Tkinter线程问题
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

# 导入项目模块
from src.modeling_clip import CLIPModel
from src.processing_clip import CLIPProcessor
from transformers.models.siglip import SiglipModel, SiglipProcessor
from utils import crop_all_subimages_by_ids, crop_all_siglip_subimages_by_ids

# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")

# 配置参数
MODEL_PATH = "models/siglip-so400m-patch14-224"  # 模型路径
TOP_K_VAR = 128  # 方差最大的前K个子图
TEMPERATURE = 0.07  # 温度参数
SAVE_VISUALIZATION = True  # 是否保存可视化结果
VIS_DIR = "visualizations/hidden_layers_analysis"

# 创建可视化目录
os.makedirs(VIS_DIR, exist_ok=True)


def validate_hidden_states_similarity():
    # 1. 加载模型和处理器
    print("正在加载模型和处理器...")
    model = SiglipModel.from_pretrained(MODEL_PATH)
    processor = SiglipProcessor.from_pretrained(MODEL_PATH)
    model.to(device)
    model.eval()
    print("模型加载完成!")
    
    # 2. 生成测试图像 - 全白图片和全黑图片
    print("正在生成测试图像: 全白图片和全黑图片")
    
    # 生成全白图片
    white_image = generate_white_image(224, 224)
    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    import requests
    black_image = Image.open(requests.get(url, stream=True).raw)
    # 生成全黑图片
    # black_image = generate_black_image(224, 224)
    
    # 3. 对每种图像进行隐藏层相似度分析
    print("\n=== 分析全白图片隐藏层相似度 ===")
    white_layer_similarities = analyze_image_hidden_states(
        model, processor, white_image, device, TOP_K_VAR, TEMPERATURE
    )
    
    print("\n=== 分析全黑图片隐藏层相似度 ===")
    black_layer_similarities = analyze_image_hidden_states(
        model, processor, black_image, device, TOP_K_VAR, TEMPERATURE
    )
    
    # 4. 可视化和保存结果
    if SAVE_VISUALIZATION:
        visualize_comparison_results(
            list(range(len(white_layer_similarities))), 
            white_layer_similarities, 
            black_layer_similarities
        )
    
    print("隐藏层相似度验证完成!")


def generate_white_image(width, height):
    """生成指定尺寸的全白图片"""
    white_array = np.full((height, width, 3), 255, dtype=np.uint8)
    white_image = Image.fromarray(white_array)
    return white_image


def generate_black_image(width, height):
    """生成指定尺寸的全黑图片"""
    black_array = np.zeros((height, width, 3), dtype=np.uint8)
    black_image = Image.fromarray(black_array)
    return black_image


def analyze_image_hidden_states(model, processor, image, device, top_k_var, temperature):
    """分析单个图像的所有隐藏层相似度"""
    # 预处理图像
    print("正在预处理图像...")
    inputs = processor(images=image, return_tensors="pt").to(device)
    
    # 直接使用vision_model获取隐藏层状态
    print("正在获取模型隐藏层状态...")
    with torch.no_grad():
        # 直接调用vision_model获取完整输出，包括隐藏层状态
        vision_outputs = model.vision_model(
            pixel_values=inputs['pixel_values'],
            output_hidden_states=True  # 确保输出隐藏层状态
        )
        vision_hidden_states = vision_outputs.hidden_states  # 获取所有隐藏层状态
    
    print(f"成功获取 {len(vision_hidden_states)} 个隐藏层状态")
    
    # 遍历每一层并计算相似度
    layer_similarities = []
    layer_indices = list(range(len(vision_hidden_states)))
    
    for layer_idx in tqdm(layer_indices, desc="处理隐藏层", position=0, leave=True):
        # 获取当前层的隐藏状态
        hidden_state = vision_hidden_states[layer_idx]
        # 提取视觉投影层
        visual_projection = model.vision_model.post_layernorm
        
        # 计算该层的特征相似度
        layer_similarity = process_single_layer(
            model, processor, image, hidden_state, visual_projection, 
            device, top_k_var, temperature
        )
        
        layer_similarities.append(layer_similarity)
        print(f"第 {layer_idx} 层 | 平均相似度: {layer_similarity:.4f}")
    
    return layer_similarities


def process_single_layer(model, processor, image, hidden_state, visual_projection, 
                         device, top_k_var, temperature):
    """处理单个隐藏层并计算相似度"""
    # 获取原始特征嵌入
    origin_local_embeds = visual_projection(hidden_state[0, :, :])  # 忽略CLS标记
    
    # 计算方差并选择方差最大的子图
    variances = torch.var(origin_local_embeds, dim=1)
    _, top_indices = torch.topk(variances, min(top_k_var, len(variances)))
    top_indices = top_indices.tolist()
    origin_local_embeds_cropped = origin_local_embeds[top_indices]
    
    # 裁剪实际子图
    try:
        local_images = crop_all_siglip_subimages_by_ids(image, 27, 27, top_indices)
    except Exception as e:
        print(f"裁剪子图失败: {e}")
        return 0.0
    
    # 获取子图的特征嵌入
    local_inputs = processor(images=local_images, return_tensors="pt").to(device)
    with torch.no_grad():
        # 使用get_image_features获取子图的特征嵌入
        target_local_embeds = model.get_image_features(**local_inputs)
    
    # 归一化特征
    origin_local_embeds_cropped = F.normalize(origin_local_embeds_cropped, dim=-1)
    target_local_embeds = F.normalize(target_local_embeds, dim=-1)
    
    # 计算余弦相似度矩阵
    cos_sim = F.cosine_similarity(
        origin_local_embeds_cropped.unsqueeze(1), 
        target_local_embeds.unsqueeze(0), 
        dim=-1
    )
    
    # 应用温度缩放
    cos_sim = cos_sim / temperature
    
    # 计算对角线元素的平均相似度（正确匹配的相似度）
    batch_size = cos_sim.size(0)
    correct_matches = [cos_sim[i, i].item() for i in range(batch_size)]
    mean_similarity = np.mean(correct_matches)
    
    # 清理变量以节省内存
    del origin_local_embeds, variances, top_indices, origin_local_embeds_cropped
    del local_images, local_inputs, target_local_embeds, cos_sim
    
    return mean_similarity


def visualize_comparison_results(layer_indices, white_layer_similarities, black_layer_similarities):
    """可视化比较两种图像的相似度结果"""
    # 创建图表
    plt.figure(figsize=(12, 6))
    plt.plot(layer_indices, white_layer_similarities, 'o-', color='blue', linewidth=2, markersize=8, label='Average Similarity of White Image')
    plt.plot(layer_indices, black_layer_similarities, 's-', color='red', linewidth=2, markersize=8, label='Average Similarity of Black Image')
    plt.title('Visual Hidden Layer Feature Similarity Analysis (White vs Black Image)', fontsize=16)
    plt.xlabel('Hidden Layer Index', fontsize=14)
    plt.ylabel('Average Feature Similarity', fontsize=14)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.legend(fontsize=12)
    plt.tight_layout()
    
    # 保存图表
    plt.savefig(os.path.join(VIS_DIR, 'hidden_layers_similarity_comparison.png'), dpi=300)
    plt.close()
    
    # 保存数值结果到CSV
    try:
        import pandas as pd
        results_df = pd.DataFrame({
            'Hidden Layer Index': layer_indices,
            'White Image Similarity': white_layer_similarities,
            'Black Image Similarity': black_layer_similarities
        })
        results_df.to_csv(os.path.join(VIS_DIR, 'hidden_layers_similarity_results.csv'), index=False)
        print(f"结果已保存到: {VIS_DIR}")
    except ImportError:
        print("pandas未安装，无法保存CSV文件")


if __name__ == "__main__":
    validate_hidden_states_similarity()